//! Custom memory allocation testing framework in Rust.
//! This module provides a safe and idiomatic Rust implementation of a memory
//! allocation testing framework, ensuring memory safety and maintaining data
//! structure invariants.

use std::alloc::{alloc, dealloc, realloc, Layout};
use std::ptr;
use std::mem;
use std::slice;

/// Magic number used to identify allocated blocks.
const ALLOC_TEST_MAGIC: u32 = 0x72ec82d2;

/// Pattern written to memory after allocation to detect uninitialized usage.
const MALLOC_PATTERN: u32 = 0xBAADF00D;

/// Pattern written to memory after deallocation to detect use-after-free.
const FREE_PATTERN: u32 = 0xDEADBEEF;

/// Header structure preceding each allocated block.
#[repr(C)]
struct BlockHeader {
    magic_number: u32,
    bytes: usize,
}

/// Count of the current number of allocated bytes.
static mut allocated_bytes: usize = 0;

/// Limit on the number of allocations possible.
static mut allocation_limit: i32 = -1;

/// Retrieves the block header for a given allocated pointer.
///
/// # Safety
/// This function assumes the pointer is valid and has a preceding `BlockHeader`.
unsafe fn alloc_test_get_header(ptr: *mut u8) -> *mut BlockHeader {
    let header_ptr = (ptr as *mut BlockHeader).offset(-1);
    assert_eq!((*header_ptr).magic_number, ALLOC_TEST_MAGIC);
    header_ptr
}

/// Overwrites a block of memory with a repeated pattern.
fn alloc_test_overwrite(ptr: *mut u8, length: usize, pattern: u32) {
    let byte_slice = unsafe { slice::from_raw_parts_mut(ptr, length) };
    for (i, byte) in byte_slice.iter_mut().enumerate() {
        let pattern_seq = (i & 3) as u32;
        *byte = ((pattern >> (8 * pattern_seq)) & 0xff) as u8;
    }
}

/// Custom malloc function with additional metadata and safety checks.
///
/// # Safety
/// This function uses unsafe Rust for memory allocation and pointer manipulation.
pub unsafe fn alloc_test_malloc(bytes: usize) -> *mut u8 {
    if allocation_limit == 0 {
        return ptr::null_mut();
    }

    let header_size = mem::size_of::<BlockHeader>();
    let total_size = header_size + bytes;

    let layout = Layout::from_size_align(total_size, mem::align_of::<BlockHeader>()).unwrap();
    let header_ptr = alloc(layout) as *mut BlockHeader;

    if header_ptr.is_null() {
        return ptr::null_mut();
    }

    (*header_ptr).magic_number = ALLOC_TEST_MAGIC;
    (*header_ptr).bytes = bytes;

    let ptr = (header_ptr as *mut u8).add(header_size);
    alloc_test_overwrite(ptr, bytes, MALLOC_PATTERN);

    allocated_bytes += bytes;

    if allocation_limit > 0 {
        allocation_limit -= 1;
    }

    ptr
}

/// Custom free function with additional safety checks.
///
/// # Safety
/// This function assumes the pointer is valid and has a preceding `BlockHeader`.
pub unsafe fn alloc_test_free(ptr: *mut u8) {
    if ptr.is_null() {
        return;
    }

    let header_ptr = alloc_test_get_header(ptr);
    let block_size = (*header_ptr).bytes;
    assert!(allocated_bytes >= block_size);

    alloc_test_overwrite(ptr, block_size, FREE_PATTERN);
    (*header_ptr).magic_number = 0;

    let header_size = mem::size_of::<BlockHeader>();
    let total_size = header_size + block_size;
    let layout = Layout::from_size_align(total_size, mem::align_of::<BlockHeader>()).unwrap();

    dealloc(header_ptr as *mut u8, layout);

    allocated_bytes -= block_size;
}

/// Custom realloc function with additional safety checks.
///
/// # Safety
/// This function assumes the pointer is valid and has a preceding `BlockHeader`.
pub unsafe fn alloc_test_realloc(ptr: *mut u8, bytes: usize) -> *mut u8 {
    let new_ptr = alloc_test_malloc(bytes);
    if new_ptr.is_null() {
        return ptr::null_mut();
    }

    if !ptr.is_null() {
        let header_ptr = alloc_test_get_header(ptr);
        let bytes_to_copy = (*header_ptr).bytes.min(bytes);
        ptr::copy_nonoverlapping(ptr, new_ptr, bytes_to_copy);
        alloc_test_free(ptr);
    }

    new_ptr
}

/// Custom calloc function with additional safety checks.
///
/// # Safety
/// This function uses unsafe Rust for memory allocation and pointer manipulation.
pub unsafe fn alloc_test_calloc(nmemb: usize, bytes: usize) -> *mut u8 {
    let total_bytes = nmemb * bytes;
    let ptr = alloc_test_malloc(total_bytes);
    if !ptr.is_null() {
        ptr::write_bytes(ptr, 0, total_bytes);
    }
    ptr
}

/// Custom strdup function with additional safety checks.
///
/// # Safety
/// This function uses unsafe Rust for memory allocation and pointer manipulation.
pub unsafe fn alloc_test_strdup(string: &str) -> *mut u8 {
    let len = string.len();
    let ptr = alloc_test_malloc(len + 1);
    if !ptr.is_null() {
        ptr::copy_nonoverlapping(string.as_ptr(), ptr, len);
        *(ptr.add(len)) = 0;
    }
    ptr
}

/// Sets the allocation limit.
pub fn alloc_test_set_limit(alloc_count: i32) {
    unsafe {
        allocation_limit = alloc_count;
    }
}

/// Retrieves the current number of allocated bytes.
pub fn alloc_test_get_allocated() -> usize {
    unsafe { allocated_bytes }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_alloc_free() {
        unsafe {
            let ptr = alloc_test_malloc(100);
            assert!(!ptr.is_null());
            assert_eq!(alloc_test_get_allocated(), 100);
            alloc_test_free(ptr);
            assert_eq!(alloc_test_get_allocated(), 0);
        }
    }

    #[test]
    fn test_realloc() {
        unsafe {
            let ptr = alloc_test_malloc(100);
            assert!(!ptr.is_null());
            let new_ptr = alloc_test_realloc(ptr, 200);
            assert!(!new_ptr.is_null());
            assert_eq!(alloc_test_get_allocated(), 200);
            alloc_test_free(new_ptr);
            assert_eq!(alloc_test_get_allocated(), 0);
        }
    }

    #[test]
    fn test_calloc() {
        unsafe {
            let ptr = alloc_test_calloc(10, 10);
            assert!(!ptr.is_null());
            assert_eq!(alloc_test_get_allocated(), 100);
            alloc_test_free(ptr);
            assert_eq!(alloc_test_get_allocated(), 0);
        }
    }

    #[test]
    fn test_strdup() {
        unsafe {
            let string = "hello";
            let ptr = alloc_test_strdup(string);
            assert!(!ptr.is_null());
            assert_eq!(alloc_test_get_allocated(), string.len() + 1);
            alloc_test_free(ptr);
            assert_eq!(alloc_test_get_allocated(), 0);
        }
    }
}
